# -*- coding: utf-8 -*-
# @Time     : 2020/8/22 2:01 下午
# @Author   : huangxiong
# @FileName : barra_attribution.py
# @Comment  : barra因子归因的基础函数
# @Software : PyCharm

import pandas as pd
import numpy as np

from quant_researcher.quant.project_tool.db_operator import db_conn
from quant_researcher.quant.datasource_fetch.portfolio_api import portfolio_tool
from quant_researcher.quant.project_tool.math_func.index_fitting import process
from quant_researcher.quant.project_tool.logger.my_logger import LOG
from quant_researcher.quant.datasource_fetch.fund_api.fund_holding_related import get_holding_report_info
from quant_researcher.quant.datasource_fetch.factor_api.factor_constant import \
    BARRA_FACTOR_NAME_DICT, BOND_FACTOR_NAME_DICT, BARRA_FULL_FACTOR, BOND_7_FACTOR
from quant_researcher.quant.datasource_fetch.stock_api import stock_price_related
from quant_researcher.quant.project_tool.exception import FOFServerError, ERR_S_2, FOFUserError, ERR_U_1

INDEX_COMPONENTS_TABLE = 'idx_components_weight'
STOCK_BARRA_EXPOSURE_TABLE = 'stk_barra_exposure'
STOCK_BARRA_COVARIANCE_TABLE = 'stk_barra_covariance_day'
STOCK_BARRA_RET_TABLE = 'stk_barra_factor_ret'
INDEX_BASEINFO_TABLE = 'idx_baseinfo'

FACTOR_NAME_DICT = BARRA_FACTOR_NAME_DICT.copy()
FACTOR_NAME_DICT.update(BOND_FACTOR_NAME_DICT)
BARRA_FACTOR_TYPE = 'BARRA_FYZ'


def get_stock_barra_exposure(calc_date):
    """
    计算股票的barra因子暴露

    :param calc_date: 计算日期，格式"2020-08-10"
    :return:
    """
    conn = db_conn.get_tk_factors_conn()
    stock_exposure_df = pd.read_sql(f"select * "
                                    f"from {STOCK_BARRA_EXPOSURE_TABLE} "
                                    f"where tradedate = '{calc_date}' ", conn)
    if stock_exposure_df.empty:
        LOG.error(f"未找到{calc_date}这天股票的barra因子暴露数据")
        return
    stock_exposure_df = stock_exposure_df.drop(columns=['tid', 'tradedate', 'sw_ind1',
                                                        'titime', 'tutime', 'tutimestamp'])
    stock_exposure_df = stock_exposure_df.rename(columns={'stockcode': 'stock_code'})
    return stock_exposure_df


def get_stock_barra_ret_cov_matrix(calc_date):
    """
    得到股票的barra因子收益率协方差矩阵

    :param calc_date: 计算日期，格式"2020-08-10"
    :return:
    """
    conn = db_conn.get_tk_factors_conn()
    stock_factor_cov_df = pd.read_sql(f"select * "
                                      f"from {STOCK_BARRA_COVARIANCE_TABLE} "
                                      f"where tradedate = '{calc_date}' ", conn)
    conn.close()
    if stock_factor_cov_df.empty:
        LOG.error(f"未找到{calc_date}这天股票的barra因子收益率协方差数据")
        return
    stock_factor_cov_df = stock_factor_cov_df.drop(columns=['tid', 'tradedate', 'titime', 'tutime', 'tutimestamp'])
    stock_factor_cov_df = stock_factor_cov_df.set_index('factor')
    # 数据库中协方差值乘以了1e6，这里调回去
    stock_factor_cov_df = stock_factor_cov_df / 1e6
    stock_factor_cov_df = stock_factor_cov_df.reindex(BARRA_FULL_FACTOR)[BARRA_FULL_FACTOR]
    stock_factor_cov_df = stock_factor_cov_df.fillna(0)
    stock_factor_cov_matrix = stock_factor_cov_df.values

    return stock_factor_cov_matrix


def get_factor_exposure(code_type, code_list, factor_type, data_df):
    """
    根据基金因子回归结果整理得到因子暴露数据

    :param code_type: fund or index
    :param code_list: 基金或指数列表
    :param factor_type: barra or bond
    :param data_df: 输入的因子回归结果
    :return:
    """
    if factor_type == 'barra':
        factor_list = BARRA_FULL_FACTOR
    else:
        factor_list = BOND_7_FACTOR

    if data_df.empty:
        exposure_df = pd.DataFrame({code_type: code_list})
        for factor_name in factor_list:
            exposure_df[factor_name] = 0
    else:
        exposure_df = data_df[[code_type, 'factor_code', 'factor_beta']]
        exposure_df = exposure_df.set_index(
            [code_type, 'factor_code'])['factor_beta'].unstack().reset_index()

    return exposure_df


def calc_risk(calc_df, factor_cov_array, calc_model):
    """
    计算组合的风险、边际风险贡献和总风险贡献

    :param calc_df: 计算的DataFrame数据
    :param factor_cov_array: 因子收益率相关性矩阵
    :param calc_model: s-barra因子 or b-债券因子
    :return:
    """
    weight_array = calc_df['weight'].values
    weight_array = weight_array.reshape(-1, 1)
    if calc_model == 's':
        factor_exposure_array = calc_df[BARRA_FULL_FACTOR].values.T
    else:
        factor_exposure_array = calc_df[BOND_7_FACTOR].values.T
    factor_weight = np.dot(factor_exposure_array, weight_array)

    # 组合风险
    sigma_p = np.dot(factor_weight.T, factor_cov_array)
    sigma_p = np.dot(sigma_p, factor_weight)
    sigma_p = np.sqrt(sigma_p[0, 0])

    if sigma_p == 0:
        LOG.error('组合风险为0，应该是哪里出了问题')
        mrc = trc = None
    else:
        # 边际风险贡献
        mrc = np.dot(factor_cov_array, factor_weight) / sigma_p
        # 总风险贡献
        trc = mrc * factor_weight

    return factor_weight, mrc, trc


def get_risk_decomposition(exposure, mrc, trc, factor_list, calc_type):
    """
    输出组合或基准的每个因子的风险暴露、边际风险贡献、总风险贡献，以及主动风险暴露、主动边际贡献、主动总风险贡献

    :param exposure: 因子暴露
    :param mrc: 边际风险贡献
    :param trc: 总风险贡献
    :param factor_list: 因子列表
    :param calc_type: 计算类型，portfolio or index
    :return:
    """
    res_df = pd.DataFrame({'factor_code': factor_list,
                           f'{calc_type}_exposure': exposure[:, 0],
                           f'{calc_type}_mrc': mrc[:, 0],
                           f'{calc_type}_trc': trc[:, 0]})

    return res_df


def portfolio_barra_exposure_by_regression(portfolio_id, calc_date):
    """
    组合基于基金净值回归后加权得到暴露与风险贡献

    :param portfolio_id: 组合ID
    :param calc_date: 计算日期，格式"2020-08-10"
    :return:
    """
    calc_date_t = calc_date.replace('-', '')
    # 根据组合代码找到持有的基金及权重
    tmp_df = portfolio_tool.get_portfolio_fund_weight(portfolio_id, calc_date_t)
    if tmp_df is None:
        LOG.error(f"{portfolio_id} 该组合id有问题")
        return
    fund_list = tmp_df['fund_code'].tolist()
    fund_weight_df = tmp_df.rename(columns={'fund_weight': 'weight'})
    # 将基金和指数代码组装，作为传入process函数的入参
    fund_bench_df = pd.DataFrame({'fund_code': fund_list})
    fund_bench_df['index_code'] = 'H11025'
    fund_bench_list = list(fund_bench_df.values)

    f_stock_df = process(code_class='fund', today=calc_date, daily_or_monthly='daily',
                         i_or_s='barra', i_or_s_specific=BARRA_FACTOR_TYPE,
                         test_codes=fund_bench_list, data_4_api=True, process_number=5)
    f_stock_df = get_factor_exposure('fund_code', fund_list, 'barra', f_stock_df)

    f_stock_df = fund_weight_df.merge(f_stock_df, how='left', on='fund_code')
    f_stock_df = f_stock_df.fillna(0)
    f_stock_df = f_stock_df[['fund_code', 'weight'] + BARRA_FULL_FACTOR]

    stock_factor_cov_matrix = get_stock_barra_ret_cov_matrix(calc_date)
    p_s_exposure, p_s_mrc, p_s_trc = calc_risk(f_stock_df, stock_factor_cov_matrix, calc_model='s')

    if p_s_mrc is None:
        return

    res_df = get_risk_decomposition(p_s_exposure, p_s_mrc, p_s_trc, BARRA_FULL_FACTOR, 'portfolio')

    return res_df


def index_barra_exposure_by_regression(index_code, calc_date):
    """
    指数基于价格回归后加权得到暴露与风险贡献

    :param index_code: 指数代码
    :param calc_date: 计算日期，格式"2020-08-10"
    :return:
    """
    index_weight_df = pd.DataFrame({'index_code': [index_code], 'weight': 1})
    i_stock_df = process(code_class='index', today=calc_date, daily_or_monthly='daily',
                         i_or_s='barra', i_or_s_specific=BARRA_FACTOR_TYPE,
                         test_codes=[index_code], data_4_api=True, process_number=1)
    if i_stock_df.empty:
        i_stock_df = pd.DataFrame(0, index=[index_code], columns=BARRA_FULL_FACTOR)
        i_stock_df = i_stock_df.reset_index().rename(columns={'index': 'index_code'})
    else:
        i_stock_df = get_factor_exposure('index_code', [index_code], 'barra', i_stock_df)

    i_stock_df = index_weight_df.merge(i_stock_df, how='left', on='index_code')
    i_stock_df = i_stock_df.fillna(0)
    i_stock_df = i_stock_df[['index_code', 'weight'] + BARRA_FULL_FACTOR]

    stock_factor_cov_matrix = get_stock_barra_ret_cov_matrix(calc_date)
    i_s_exposure, i_s_mrc, i_s_trc = calc_risk(i_stock_df, stock_factor_cov_matrix, calc_model='s')

    res_df = get_risk_decomposition(i_s_exposure, i_s_mrc, i_s_trc, BARRA_FULL_FACTOR, 'index')

    return res_df


def portfolio_barra_exposure_by_holding(portfolio_id, calc_date):
    """
    根据组合基金持仓得到暴露与风险贡献

    :param portfolio_id: 组合ID
    :param calc_date: 计算日期，格式"2020-08-10"
    :return:
    """
    # 根据组合代码找到持有的基金及权重
    calc_date_t = calc_date.replace('-', '')
    tmp_df = portfolio_tool.get_portfolio_fund_weight(portfolio_id, calc_date_t)
    if tmp_df is None:
        LOG.error(f"{portfolio_id} 该组合id有问题")
        return
    fund_list = tmp_df['fund_code'].tolist()

    # 去找基金最近的股票持仓数据
    fund_hold_df = get_holding_report_info(report='stock', fund_code=fund_list, report_type=[5, 6],
                                           end_date=calc_date, select=['fund_code', 'stock_code', 'hold_mktvalue as value',
                                                                       'pptinnv as stock_weight'])
    if fund_hold_df is None:
        return
    fund_hold_df['stock_code'] = fund_hold_df['stock_code'].str.extract(r'(\d+)')
    fund_hold_df['stock_code'] = fund_hold_df['stock_code'].str.zfill(5)
    fund_hold_df = fund_hold_df.dropna()
    fund_hold_df['stock_weight'] = fund_hold_df['stock_weight'] / 100
    stock_exposure_df = get_stock_barra_exposure(calc_date)
    if stock_exposure_df is None:
        return
    f_stock_df = fund_hold_df.groupby('stock_code')['value'].sum().reset_index()
    f_stock_df['weight'] = f_stock_df['value'] / f_stock_df['value'].sum()
    f_stock_df = f_stock_df.drop(columns='value')
    f_stock_df = f_stock_df.merge(stock_exposure_df, how='left', on='stock_code')
    f_stock_df = f_stock_df.fillna(0)

    stock_factor_cov_matrix = get_stock_barra_ret_cov_matrix(calc_date)
    p_s_exposure, p_s_mrc, p_s_trc = calc_risk(f_stock_df, stock_factor_cov_matrix, calc_model='s')

    res_df = get_risk_decomposition(p_s_exposure, p_s_mrc, p_s_trc, BARRA_FULL_FACTOR, 'portfolio')

    return res_df


def index_barra_exposure_by_holding(index_code, calc_date):
    """
    指数基于持仓得到暴露与风险贡献

    :param index_code: 指数代码
    :param calc_date: 计算日期，格式"2020-08-10"
    :return:
    """
    calc_date_t = calc_date.replace('-', '')
    conn = db_conn.get_basic_data_conn()
    # 找到指数在离calc_date这天最近的成份股及权重
    index_latest_weight_date = pd.read_sql(f"select max(trade_date) "
                                           f"from {INDEX_COMPONENTS_TABLE} "
                                           f"where index_code = '{index_code}' "
                                           f"and trade_date <= '{calc_date_t}' ", conn).iloc[0, 0]

    # 找到所有债券指数的列表
    bond_index_list = pd.read_sql(f"select index_code "
                                  f"from {INDEX_BASEINFO_TABLE} "
                                  f"where index_type = '03' ", conn)['index_code'].tolist()

    # 如果能找到指数的成份股，且指数不是债券指数，则用成份股计算
    if (index_latest_weight_date is not None) and (index_code not in bond_index_list):
        index_weight_df = pd.read_sql(f"select secu_code as stock_code, weight/100 as weight "
                                      f"from {INDEX_COMPONENTS_TABLE} "
                                      f"where index_code = '{index_code}' "
                                      f"and trade_date = '{index_latest_weight_date}' ", conn)
        # 如果指数权重最新日期比现在日期小，按照股票涨跌幅调整个股权重
        if index_latest_weight_date < calc_date_t:
            # stock_ret = common_functions.get_stock_total_ret_between_days(index_latest_weight_date, calc_date_t)
            stock_ret = stock_price_related.get_stock_total_return(index_latest_weight_date, calc_date_t)
            index_weight_df = index_weight_df.merge(stock_ret, how='left', on='stock_code').fillna(0)
            index_weight_df['weight2'] = index_weight_df['weight'] * (1 + index_weight_df['ret'])
            index_weight_df['weight_adj'] = index_weight_df['weight2'] / index_weight_df['weight2'].sum()
            index_weight_df = index_weight_df[
                ['stock_code', 'weight_adj']].rename(columns={'weight_adj': 'weight'})

        stock_exposure_df = get_stock_barra_exposure(calc_date)
        i_stock_df = index_weight_df.merge(stock_exposure_df, how='left', on='stock_code')
        if i_stock_df.dropna().empty:
            LOG.error(f"{index_code}指数的成份股没有股票，可能是债券指数，且在指数信息表中没有分类")
            conn.close()
            return
        i_stock_df = i_stock_df.fillna(0)
        stock_factor_cov_matrix = get_stock_barra_ret_cov_matrix(calc_date)
        i_s_exposure, i_s_mrc, i_s_trc = calc_risk(i_stock_df, stock_factor_cov_matrix, calc_model='s')

        res_df = get_risk_decomposition(i_s_exposure, i_s_mrc, i_s_trc, BARRA_FULL_FACTOR, 'index')
        return res_df
    else:
        LOG.error(f"没有找到{index_code}指数的成份股信息")
        conn.close()
        return


def get_portfolio_index_barra_attribution(portfolio_id, calc_date, customer_create, index_code, customer_index):
    """
    得到组合基于回归的barra因子暴露，和组合和基准基于持仓的barra因子暴露

    :param portfolio_id: 组合ID
    :param calc_date: 计算日期，格式"2020-08-10"
    :param customer_create: 是否是用户自建指数，1-是，0-否，默认0
    :param index_code: 业绩基准代码
    :param customer_index: 用户自建指数，customer_create=1时必须传入该参数。
                           需要传入构建自建指数所需要的指数代码及权重，格式{"CBA00301.CS": 0.5, "000300": 0.5}
    :return:
    """
    factor_name_df = pd.DataFrame(list(FACTOR_NAME_DICT.items()), columns=['factor_code', 'factor_name'])

    p_attribution_by_holding = portfolio_barra_exposure_by_holding(portfolio_id, calc_date)
    if p_attribution_by_holding is None:
        raise FOFServerError(ERR_S_2, f"未找到该组合 {portfolio_id} 的持仓信息")

    if (customer_create == 1 and customer_index is None) or \
            (customer_create == 0 and index_code is None):
        raise FOFUserError(ERR_U_1, "如果customer_create为1，需要输入customer_index，"
                                    "如果customer_create为0，需要输入index_code，"
                                    "请按要求重新输入")
    if customer_create == 0:
        b_attribution_by_holding = index_barra_exposure_by_holding(index_code, calc_date)
        if b_attribution_by_holding is None:
            raise FOFServerError(ERR_S_2, f"未找到基准 {index_code} 的成分股信息")
    else:
        tmp_list = []
        for i in customer_index.keys():
            tmp_df = index_barra_exposure_by_holding(i, calc_date)
            if tmp_df is not None:
                tmp_df[['index_exposure', 'index_mrc', 'index_trc']] = \
                    tmp_df[['index_exposure', 'index_mrc', 'index_trc']] * customer_index[i]
                tmp_list.append(tmp_df)
        if len(tmp_list) == 0:
            raise FOFServerError(ERR_S_2, f"未找到基准 {customer_index} 的成分股信息")
        b_attribution_by_holding = pd.concat(tmp_list)
        b_attribution_by_holding = b_attribution_by_holding.groupby('factor_code')[['index_exposure', 'index_mrc', 'index_trc']].sum().reset_index()

    res_df = p_attribution_by_holding.merge(b_attribution_by_holding, how='outer', on='factor_code')
    res_df = res_df.set_index('factor_code')
    res_df['active_exposure'] = res_df['portfolio_exposure'] - res_df['index_exposure']
    res_df['active_mrc'] = res_df['portfolio_mrc'] - res_df['index_mrc']
    res_df['active_trc'] = res_df['portfolio_trc'] - res_df['index_trc']
    # barra因子收益率
    conn = db_conn.get_tk_factors_conn()
    barra_ret = pd.read_sql(f"select * "
                            f"from {STOCK_BARRA_RET_TABLE} "
                            f"where tradedate = '{calc_date}' ", conn)
    conn.close()
    if barra_ret.empty:
        raise FOFServerError(ERR_S_2, f"没有找到{calc_date}这天的barra因子收益率数据")
    barra_ret = barra_ret.drop(columns=['tid', 'tradedate', 'r2', 'titime', 'tutime', 'tutimestamp'])
    barra_ret = barra_ret.T.rename(columns={0: 'ret'})
    res_df['portfolio_ret'] = res_df['portfolio_exposure'] * barra_ret['ret']
    res_df['index_ret'] = res_df['index_exposure'] * barra_ret['ret']
    res_df['active_ret'] = res_df['portfolio_ret'] - res_df['index_ret']
    res_df = res_df.fillna(0)
    res_df = res_df.drop(columns=['portfolio_mrc', 'index_mrc', 'active_mrc'])
    res_df = res_df.round(4)
    res_df = res_df.reset_index()

    res_df = res_df.merge(factor_name_df, how='left', on='factor_code')
    res_list = res_df.to_dict(orient='records')

    return res_list


if __name__ == '__main__':
    # portfolio_barra_exposure_by_regression('1189', '20200810')
    aaa = get_portfolio_index_barra_attribution('1352', '2021-08-10', 0, '000300', None)
